///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file conjugate_gradient.h
///
/// \brief   header file for the conjugate gradient solver
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_CONJUGATE_GRADIENT_HEADER
#define PCPS_CONJUGATE_GRADIENT_HEADER

#include <src/pcps.h>
#include <src/matrix_actor.h>

namespace pcps {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Solves a system of linear equations A x = b using the conjugate gradient method.
  ///         
  ///         Using preconditioner M, the equations become
  ///         
  ///         (M A M) v = (M b)           v = M^(-1) x
  ///         
  ///         which is solved for v by conjugate gradient.
  ///         
  ///         x is then computed as x = M v.
  ///         
  ///         A must be a positive definite hermitian matrix.
  ///         M must be an invertable hermitian matrix.
  ///         
  ///         Presumably, the ratio of largest to smallest
  ///         eigenvalues is smaller for (M A M) than for A.
  ///
  /// \param[in]      n        the dimension of the matrices and vectors
  /// \param[in]      m        the maximum number of iterations
  /// \param[in]      thresh   threshold for convergence
  /// \param[out]     x        on exit, the solution vector x.
  /// \param[in,out]  b        on entry, the vector b.    on exit, contents are overwritten.
  /// \param[out]     y        vector used as workspace.  on exit, contents are overwritten.
  /// \param[out]     d        vector used as workspace.  on exit, contents are overwritten.
  /// \param[out]     r        vector used as workspace.  on exit, contents are overwritten.
  /// \param[in,out]  A        an object that multiplies by the matrix A
  /// \param[in,out]  M        an object that multiplies by the matrix M
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template<class S> inline void conjugate_gradient(const int n,
                                                   const int m,
                                                   const double thresh,
                                                   S * const x,
                                                   S * const b,
                                                   S * const y,
                                                   S * const d,
                                                   S * const r,
                                                   pcps::MatrixActor<S> * A,
                                                   pcps::MatrixActor<S> * M) {

    S r_norm = pcps::unity<S>() * 1.0e100;
    S alpha, beta;

    // initialize the residual as r = M b
    M->operate(n, b, r);

    // initialize the search direction as d = r
    pcps::xcopy(n, r, 1, d, 1);

    // use the b vector to hold v
    S * const v = b;

    // initialize v to zero
    pcps::xscal(n, pcps::zero<S>(), v, 1);

    // compute the initial residual norm
    r_norm = pcps::xdotc(n, r, 1, r, 1);

    // iteratively compute v
    int iter = 0;
    for ( ; iter < m; iter++) {

      // check if we are finished  (only the result on the root process is meaningful)
      bool converged = ( std::abs(std::sqrt(r_norm)) < thresh );
      pcps::bcast(converged);
      if (converged) break;

      // compute y = (M A M) d       (use x as an intermediate)
      M->operate(n, d, y);
      A->operate(n, y, x);
      M->operate(n, x, y);

      // compute alpha
      alpha = r_norm / pcps::xdotc(n, d, 1, y, 1);

      // update v
      pcps::xaxpy(n, alpha, d, 1, v, 1);

      // update the residual
      pcps::xaxpy(n, -alpha, y, 1, r, 1);

      // compute beta
      beta = pcps::xdotc(n, r, 1, r, 1) / r_norm;

      // compute new residual norm
      r_norm = beta * r_norm;

      // compute new search direction
      pcps::xscal(n, beta, d, 1);
      pcps::xaxpy(n, pcps::unity<S>(), r, 1, d, 1);

    }

    // check if the method converged
    const bool converged = ( std::abs(std::sqrt(r_norm)) < thresh );
    if (converged && pcps::mpi::rank() == 0)
      std::cout << boost::format("conjugate gradient converged in %10i iterations") % iter << std::endl << std::endl;
    else if (pcps::mpi::rank() == 0)
      std::cout << boost::format("conjugate gradient did not converge after %10i iterations") % iter << std::endl << std::endl;
      //throw pcps::Exception( (boost::format("pcps::conjugate_gradient did not converge.  residual = %.2e") % std::abs(std::sqrt(r_norm))).str() );

    // we currently have the solution to the equation (M A M) (M^(-1) x) = (M b), so multiply by M to get solution to A x = b
    A->operate(n, v, x);

  }

} // end namespace pcps

#endif
